#include <exception>
#include <map>
#include <algorithm>
#include "Core/Wrapper_LAPACKE.h"
#include "Core/VecCompare.h"
#include "Embedded/BlockRelaxation.h"
#include "../example/common/TestUtility.h"

template <int Dim>
BlockRelaxation<Dim>::BlockRelaxation(const EmGrid &emGrid, const RowSparse<iVec> &aSparseOp, Real splitThr)
    : originalOp(aSparseOp)
{
  const auto dx = emGrid.spacing();
  VecCompare<int, Dim> ivcmp;
  std::map<iVec, int, VecCompare<int, Dim>> unknownLookup;
  // initialize the unknown lookup
  auto rit = originalOp.cbegin();
  int n;
  std::size_t nn;
  for(nn = 0; nn < originalOp.getNumRows(); ++nn) {
    unknownLookup.insert(std::make_pair(rit.getRow(), nn));
    ++rit;
  }
  n = nn;
  // enlarge the unknown block
  RowSparse<iVec> extendedOp = originalOp;
  for(rit = originalOp.cbegin(); rit != originalOp.cend(); ++rit) {
    // find the diagonal term
    auto sbegin = &rit.getColumn(0);
    auto send = &rit.getColumn(rit.getNz());
    auto jit = std::find_if(sbegin, send, [&](const auto &rhs) { return ivcmp.compare(rit.getRow(), rhs) == 0; });
    assert(jit != send);
    Real diagTerm = std::abs(rit.getValue(jit - sbegin));
    // check any column with large coefficient
    for(int m = 0; m < rit.getNz(); ++m) {
      if(m == jit - sbegin)
        continue;
      if(std::abs(rit.getValue(m)) < splitThr * diagTerm)
        continue;
      auto extraRow = rit.getColumn(m);
      if(emGrid.nodeLabels(extraRow) < 0 || norm(extraRow - rit.getRow(), 0) > 1)
        continue;
      auto result = unknownLookup.insert(std::make_pair(extraRow, n));
      if(!result.second)
        continue;
      ++n;
      // recover the interior difference scheme
      int i = extraRow[0], j = extraRow[1], k = (Dim>=3) ? (extraRow[2]) : (0);
      iVec tmpCols[] = { {i, j, k},
                         {i-2, j, k}, {i-1, j, k}, {i+1, j, k}, {i+2, j, k},
                         {i, j-2, k}, {i, j-1, k}, {i, j+1, k}, {i, j+2, k},
                         {i, j, k-2}, {i, j, k-1}, {i, j, k+1}, {i, j, k+2} };
      Real tmpVals[] = { -30.0/12 * (1.0/(dx[0]*dx[0]) + 1.0/(dx[1]*dx[1]) + (Dim>=3) * 1.0/(dx[2]*dx[2])),
                         -1.0/12/(dx[0]*dx[0]), 16.0/12/(dx[0]*dx[0]), 16.0/12/(dx[0]*dx[0]), -1.0/12/(dx[0]*dx[0]),
                         -1.0/12/(dx[1]*dx[1]), 16.0/12/(dx[1]*dx[1]), 16.0/12/(dx[1]*dx[1]), -1.0/12/(dx[1]*dx[1]),
                         0.0, 0.0, 0.0, 0.0 };
      if(Dim >= 3) {
        tmpVals[9] = tmpVals[12] = -1.0/12/(dx[2]*dx[2]);
        tmpVals[10] = tmpVals[11] = 16.0/12/(dx[2]*dx[2]);
      }
      extendedOp.insertRow(extraRow, (Dim>=3) ? (13) : (9), tmpCols, tmpVals);
      dbgcout2 << "Include row " << extraRow << " in the diagonal block. \n";
    }
  }
  blockSize = n;
  dbgcout1 << "BlockRelaxation : blockSize = " << blockSize << std::endl;
  // split the irregular part of the Laplacian matrix into the diagonal block and the off-diagonal block
  diagBlock.resize(Vec<int, 2> {n, n});
  diagBlock = 0.0;
  for(rit = extendedOp.cbegin(); rit != extendedOp.cend(); ++rit) {
    int nz = rit.getNz();
    assert(unknownLookup.count(rit.getRow()) > 0);
    int i = unknownLookup[rit.getRow()];
    int j;
    std::vector<iVec> offDiagCols;
    std::vector<Real> offDiagVals;
    for(int m = 0; m < nz; ++m) {
      auto result = unknownLookup.find(rit.getColumn(m));
      if(result == unknownLookup.cend()) { // to the off-diagonal block
        offDiagCols.push_back(rit.getColumn(m));
        offDiagVals.push_back(rit.getValue(m));
      } else { // to the diagonal block
        j = result->second;
        diagBlock(i, j) = rit.getValue(m);
      }
    }
    offDiagBlock.insertRow(rit.getRow(), offDiagCols.size(), offDiagCols.data(), offDiagVals.data());
  }
  // pre-cache the LU factorization
  ipiv.resize(blockSize);
  CPUTimer timer;
  auto info = LAPACKE_dgetrf(LAPACK_COL_MAJOR, blockSize/*m*/, blockSize/*n*/,
                             diagBlock.data()/*a*/, blockSize/*lda*/, ipiv.data()/*ipiv*/);
  std::cout << "Factorization : " << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
  if(info != 0)
    throw std::runtime_error("[BlockRelaxation::BlockRelaxation] DGETRF fails.");
}

template<int Dim>
void BlockRelaxation<Dim>::apply(Tensor<Real, Dim> &aPhi, const Tensor<Real, Dim> &rhs) const
{
  if(blockSize == 0)
    return;
  Tensor<Real, 1> rr = offDiagBlock.AXPY(-1.0, aPhi, rhs);
  auto info = LAPACKE_dgetrs(LAPACK_COL_MAJOR, 'N'/*trans*/, blockSize/*n*/, 1/*nrhs*/,
                             diagBlock.data()/*a*/, blockSize/*lda*/, ipiv.data()/*ipiv*/,
                             rr.data()/*b*/, blockSize/*ldb*/);
  if(info != 0)
    throw std::runtime_error("[BlockRelaxation:apply] DGETRS fails. ");
  int m = 0;
  for(auto rit = offDiagBlock.cbegin(); rit != offDiagBlock.cend(); ++rit, ++m)
    aPhi(rit.getRow()) = rr(m); //aPhi(rit.getRow()) * (1-weightForJacobi) + rr(m) * weightForJacobi;
}

template class BlockRelaxation<SpaceDim>;
